Federated Learning - SMS spam prediction with a GRU model

In this notebook, we will train a model using federated approach.

NOTE: At the time of running this notebook, we were running the PyGrid components in background mode.

NOTE: Components:

To start the network:

  • git clone https://github.com/OpenMined/PyGridNetwork
  • cd PyGridNetwork
  • python -m gridnetwork --start_local_db --port=5000

To start one grid node:

  • git clone https://github.com/OpenMined/PyGridNode

  • python -m gridnode --start_local_db --id=alice --port=3001 --gateway_url=http://localhost:5000

This notebook was made based on Federated SMS Spam prediction.

Authors:

Useful imports


In [1]:
import numpy as np

import syft as sy
from syft.grid.public_grid import PublicGridNetwork

import torch as th
from torch import optim

import warnings

warnings.filterwarnings("ignore")

Connecting to Grid Network


In [2]:
hook = sy.TorchHook(th)

my_grid = PublicGridNetwork(hook, "http://localhost:5000")

Seach a dataset


In [3]:
data = my_grid.search("#X", "#spam", "#dataset")
target = my_grid.search("#Y", "#spam", "#dataset")

In [4]:
data


Out[4]:
{'Bob': [(Wrapper)>[PointerTensor | me:96906351671 -> Bob:58176119719]
  	Tags: #spam #X #dataset 
  	Shape: torch.Size([2786, 30])
  	Description: The input datapoints to the SPAM dataset....],
 'Alice': [(Wrapper)>[PointerTensor | me:91520735976 -> Alice:16627579692]
  	Tags: #spam #X #dataset 
  	Shape: torch.Size([2786, 30])
  	Description: The input datapoints to the SPAM dataset....]}

In [5]:
target


Out[5]:
{'Bob': [(Wrapper)>[PointerTensor | me:52053241859 -> Bob:34347751563]
  	Tags: #Y #spam #dataset 
  	Shape: torch.Size([2786])
  	Description: The input labels to the SPAM dataset....],
 'Alice': [(Wrapper)>[PointerTensor | me:44687882238 -> Alice:89030703951]
  	Tags: #Y #spam #dataset 
  	Shape: torch.Size([2786])
  	Description: The input labels to the SPAM dataset....]}

Load the model


In [6]:
from handcrafted_GRU import GRU

Parameters


In [7]:
data = list(data.values())
target = list(target.values())

VOCAB_SIZE = 0
for data_comp in data:
    VOCAB_SIZE = max(VOCAB_SIZE, int(data_comp[0].max().get()))
    
VOCAB_SIZE += 1
HIDDEN_DIM = 10
EMBEDDING_DIM = 50
BATCH_SIZE = 128
CLIP = 5 # gradient clipping - to avoid gradient explosion (frequent in RNNs)
DROPOUT = 0.2
EPOCHS = 15
LR = 0.1

In [8]:
# Initiating the model
model = GRU(vocab_size=VOCAB_SIZE, hidden_dim=HIDDEN_DIM, embedding_dim=EMBEDDING_DIM, dropout=DROPOUT)

# And the optimizer
optimizer = optim.SGD(model.parameters(), lr=LR)

# And the loss
criterion = th.nn.BCELoss()

Perform train


In [ ]:
import math # Needed for separating into batches

def train(epoch):
    dataset_size = sum([len(data[i][0]) for i in range(len(data))])
    model.train()
    
    for i in range(len(data)):
        loss_cum = 0
        nr_batches = math.ceil(len(data[i][0]) / BATCH_SIZE)
        for batch_idx in range(nr_batches):
            # Extract the batch for training and target
            data_batch = data[i][0][BATCH_SIZE * batch_idx : BATCH_SIZE * (batch_idx + 1), :]
            target_batch = target[i][0][BATCH_SIZE * batch_idx : BATCH_SIZE * (batch_idx + 1)]
            
            # Send the model to the worker
            worker = data_batch.location
            model.send(worker)
            h = th.Tensor(np.zeros((data_batch.shape[0], HIDDEN_DIM))).send(worker)
            
            optimizer.zero_grad()
            pred, _ = model(data_batch, h)
            loss = criterion(pred.squeeze(), target_batch.float())
            loss.backward()
            optimizer.step()
            model.get()
            
            # Cumulate the loss
            loss_cum += loss.get().item()
        
        print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, i * data[i][0].shape[0], dataset_size,
                       100. *  (i * data[i][0].shape[0]) / dataset_size, loss_cum))

for epoch in range(EPOCHS):
    train(epoch)


Train Epoch: 0 [0/5572 (0%)]	Loss: 12.313264
Train Epoch: 0 [2786/5572 (50%)]	Loss: 9.950761
Train Epoch: 1 [0/5572 (0%)]	Loss: 9.132653